import cv2
import numpy as np
import tensorflow as tf
import os
from hitnet_core import HitNet, ModelType

data_path = "/home/whitby/Documents/data/kitti/02/"
stereo_bf = 386.1448    # b * f

model_type = ModelType.middle_bury
model_path = "models/middlebury_d400.pb"
hitnet_depth = HitNet(model_path, stereo_bf, model_type)

image_index = 0

while True:
    left_path = data_path + "image_0/%06d.png" % image_index
    right_path = data_path + "image_1/%06d.png" % image_index
    if not os.path.isfile(left_path) or not os.path.isfile(right_path):
        print("image file is not exist")
        break

    left_img = cv2.imread(left_path, cv2.IMREAD_COLOR)
    right_img = cv2.imread(right_path, cv2.IMREAD_COLOR)

    disparity_map = hitnet_depth(left_img, right_img)
    depth_map = hitnet_depth.get_depth()

    # np.savetxt(data_path + "depth_0/%06d.txt" % image_index, depth_map, fmt='%.3f')
    cv2.imwrite(data_path + "depth_0/%06d.tiff" % image_index, depth_map)
    # np.save(data_path + "depth_0/%06d.npy" % image_index, depth_map)

    image_index += 1

    if image_index % 100 == 0:
        print(image_index)
